#include <iostream>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <functional>
#include <queue>
#include <vector>
using namespace std;

class ThreadPool {
public:
    typedef function<void()> Task;

    ThreadPool();
    void setMaxTaskQueueSize(int size);
    void start(int threadNum);
    void stop();
    void pushTask(Task fn);

private:
    vector<thread> threads;
    queue<Task> tasks;
    mutex qmtx;
    condition_variable cv;
    bool stopFlag;
    size_t maxTaskQueueSize;
};

ThreadPool::ThreadPool() {
    stopFlag = false;
    maxTaskQueueSize = 0;
}

void ThreadPool::setMaxTaskQueueSize(int size) {
    maxTaskQueueSize = size;
}

void ThreadPool::start(int threadNum) {
    for (int i = 0; i < threadNum; i++) {
        threads.emplace_back([this]() {
            while (true) {
                Task task;
                {
                    unique_lock<mutex> lock(qmtx);
                    cv.wait(lock, [this]() { return stopFlag || !tasks.empty(); });
                    if (stopFlag && tasks.empty()) {
                        return;
                    }
                    task = move(tasks.front());
                    tasks.pop();
                }
                task();
            }
            });
    }
}

void ThreadPool::stop() {
    {
        unique_lock<mutex> lock(qmtx);
        stopFlag = true;
    }
    cv.notify_all();

    for (auto& worker : threads) {
        worker.join();
    }
}

void ThreadPool::pushTask(Task fn) {
    {
        unique_lock<mutex> lock(qmtx);
        cv.wait(lock, [this]() { return tasks.size() < maxTaskQueueSize; });
        tasks.push(fn);
    }
    cv.notify_one();
}

int main() {
    ThreadPool pool;
    pool.start(10); // 将最大线程数设置为10
    pool.setMaxTaskQueueSize(25);

    // 测试用例 1: 推入 5 个任务
    cout << "测试用例 1: 推入 5 个任务" << endl;
    for (int i = 0; i < 5; i++) {
        pool.pushTask([i]() {
            static mutex iomtx;
            {
                lock_guard<mutex> lock(iomtx);
                cout << "任务 " << i << " 执行" << endl;
            }
            this_thread::sleep_for(chrono::seconds(1));
            });
    }
    this_thread::sleep_for(chrono::seconds(6)); // 等待任务执行完成

    // 测试用例 2: 推入 20 个任务
    cout << "测试用例 2: 推入 20 个任务" << endl;
    for (int i = 0; i < 20; i++) {
        pool.pushTask([i]() {
            static mutex iomtx;
            {
                lock_guard<mutex> lock(iomtx);
                cout << "任务 " << i << " 执行" << endl;
            }
            this_thread::sleep_for(chrono::seconds(1));
            });
    }
    this_thread::sleep_for(chrono::seconds(21)); // 等待任务执行完成

    pool.stop();
    return 0;
}
